"""
网络结构
"""
import mindspore.nn as nn


class MsMLP(nn.Cell):
    def __init__(self, input_dim, hidden_dim, n_layers, output_dim):
        super(MsMLP, self).__init__()
        # 初始化输入层
        net_layer = [nn.Dense(input_dim, hidden_dim, has_bias=True, bias_init='normal'), nn.LeakyReLU(alpha=0.01)]
        # 加入隐层和输出层数
        for i in range(n_layers):
            if i != n_layers - 1:
                net_layer.append(nn.Dense(hidden_dim, hidden_dim, has_bias=True, bias_init='normal'))
                net_layer.append(nn.LeakyReLU(alpha=0.01))
            else:
                net_layer.append(nn.Dense(hidden_dim, output_dim, has_bias=True, bias_init='normal'))
        self.seq = nn.SequentialCell(net_layer)

    def construct(self, x):
        return self.seq(x)